#! /usr/bin/env python3
# coding: utf-8
"""
References:
[1]Fast Multichannel Nonnegative Matrix Factorization with Directivity-Aware Jointly-Diagonalizable
Spatial Covariance Matrices for Blind Source Separation
[2]Fast Multichannel Source Separation Based on Jointly Diagonalizable Spatial Covariance Matrices
"""

import numpy as np
import chainer
import sys, os
import librosa
import soundfile as sf
import time
import pickle as pic

from configure_FastModel import *
from FastFCA_online import FastFCA_online
from FastMNMF_online import FastMNMF_online

try:
    from chainer import cuda
    FLAG_GPU_Available = True
except:
    print("---Warning--- You cannot use GPU acceleration because chainer or cupy is not installed")


class FastMNMF2_online(FastMNMF_online):  # 继承父类
    """ Blind Source Separation Using Fast Multichannel Nonnegative Matrix Factorization (FastMNMF)

    X_FTM: the observed complex spectrogram
    Q_FMM: diagonalizer that converts a spatial covariance matrix (SCM) to a diagonal matrix
    G_NFM: diagonal elements of the diagonalized SCMs
    W_NFK: basis vectors for each source
    H_NKT: activations for each source
    lambda_NFT: power spectral densities of each source (W_NFK @ H_NKT)
    Qx_power_FTM: power spectra of Qx
    Y_FTM: \sum_n lambda_NFT G_NM
    """

    def __init__(self, n_source=2, n_basis=8, xp=np, init_SCM="unit", dl_ratio=0, total_frame=0, all_sep_spec=0, first_batch_size=7, weight_rho=0.9, mini_batch_size=4):
        """ initialize FastMNMF

        Parameters:
        -----------
            n_source: int
                the number of sources
            n_iteration: int
                the number of iteration to update all variables
            n_basis: int
                the number of bases of each source
            init_SCM: str
                how to initialize covariance matrix {unit, obs, ILRMA}
        """
        super(FastMNMF2_online, self).__init__(n_source=n_source, xp=xp, init_SCM=init_SCM, dl_ratio=dl_ratio)
        self.n_basis = n_basis
        self.method_name = "FastMNMF2_online"
        self.total_frame = total_frame
        self.all_sep_spec = all_sep_spec
        self.first_batch_size = first_batch_size
        self.weight_rho = weight_rho
        self.mini_batch_size = mini_batch_size


    def set_parameter(self, n_source=None, n_basis=None, init_SCM=None):
        """ set parameters

        Parameters:
        -----------
            n_source: int
            n_iteration: int
            n_basis: int
            init_SCM: str
                how to initialize covariance matrix {unit, obs, ILRMA}
        """
        super(FastMNMF2_online, self).set_parameter(n_source=n_source, init_SCM=init_SCM)  # super()调用父类
        if n_basis != None:
            self.n_basis = n_basis


    ## 初始化first_batch的pre_batch参数，均设为0，等效为离线
    def initialize_pre_batch(self):
        # W参数初始化
        self.pre_batch_a_W = self.xp.zeros([self.n_source, self.n_freq, self.n_basis])  # N F K
        self.pre_batch_b_W = self.xp.zeros([self.n_source, self.n_freq, self.n_basis])  # N F K
        self.pre_batch_W_NFK = self.xp.zeros([self.n_source, self.n_freq, self.n_basis])  # N F K
        # G参数初始化
        self.pre_batch_G_a_1 = self.xp.zeros([self.n_source, self.n_mic])  # N M
        self.pre_batch_G_b_1 = self.xp.zeros([self.n_source, self.n_mic])  # N M
        self.pre_batch_G_NM = self.xp.zeros([self.n_source, self.n_mic])  # N M


    def update(self):
        self.update_WH()
        self.update_CovarianceDiagElement()
        self.udpate_Diagonalizer()
        self.normalize()


    def update_WH(self):
        tmp1_NFT = (self.G_NM[:, None, None] * (self.Qx_power_FTM / (self.Y_FTM ** 2))[None]).sum(axis=3)  # numerator of (34)(35) in [1]
        tmp2_NFT = (self.G_NM[:, None, None] / self.Y_FTM[None]).sum(axis=3)  # denominator of (34)(35) in [1]
        self.a_W = (self.H_NKT[:, None] * tmp1_NFT[:, :, None]).sum(axis=3)  # N F K T -> N F K
        self.b_W = (self.H_NKT[:, None] * tmp2_NFT[:, :, None]).sum(axis=3)  # N F K T -> N F K
        # online version
        Fj_WaW_NFK = self.W_NFK * self.a_W * self.W_NFK + self.weight_rho * self.pre_batch_W_NFK * self.pre_batch_a_W * self.pre_batch_W_NFK
        Fj_b_NFK = self.b_W + self.weight_rho * self.pre_batch_b_W
        a_H = (self.W_NFK[..., None] * tmp1_NFT[:, :, None]).sum(axis=1) # N F K T -> N K T
        b_H = (self.W_NFK[..., None] * tmp2_NFT[:, :, None]).sum(axis=1) # N F K T -> N K T
        self.W_NFK = self.xp.sqrt(Fj_WaW_NFK / Fj_b_NFK)  # (34) of [1]
        self.H_NKT = self.H_NKT * self.xp.sqrt(a_H / b_H)  # (35) of [1]

        self.lambda_NFT = self.W_NFK @ self.H_NKT + EPS
        self.Y_FTM = (self.lambda_NFT[..., None] * self.G_NM[:, None, None]).sum(axis=0)


    def update_CovarianceDiagElement(self):
        self.G_a_1 = (self.lambda_NFT[..., None] * (self.Qx_power_FTM / (self.Y_FTM ** 2))[None]).sum(axis=2).sum(axis=1)  # N F T M
        self.G_b_1 = (self.lambda_NFT[..., None] / self.Y_FTM[None]).sum(axis=2).sum(axis=1)
        Fj_GaG_NM = self.G_NM * self.G_a_1 * self.G_NM + self.weight_rho * self.pre_batch_G_NM * self.pre_batch_G_a_1 * self.pre_batch_G_NM
        Fj_b_NM = self.G_b_1 + self.weight_rho * self.pre_batch_G_b_1
        self.G_NM = self.xp.sqrt(Fj_GaG_NM / Fj_b_NM)  # (36) of [1]
        self.G_NM += EPS
        self.Y_FTM = (self.lambda_NFT[..., None] * self.G_NM[:, None, None]).sum(axis=0)


    def normalize(self):
        phi_F = self.xp.trace(self.Q_FMM * self.Q_FMM.conj(), axis1=1, axis2=2).real / self.n_mic  # (37) of [1]
        self.Q_FMM = self.Q_FMM / self.xp.sqrt(phi_F)[:, None, None]
        self.W_NFK = self.W_NFK / phi_F[None, :, None]

        mu_N = (self.G_NM).sum(axis=1).real  # (38) of [1]
        self.G_NM = self.G_NM / mu_N[:, None]
        self.W_NFK = self.W_NFK * mu_N[:, None, None]

        nu_NK = self.W_NFK.sum(axis=1)  # (28) of [1]
        self.W_NFK = self.W_NFK / nu_NK[:, None]
        self.H_NKT = self.H_NKT * nu_NK[:, :, None]
        self.lambda_NFT = self.W_NFK @ self.H_NKT + EPS

        self.reset_variable()


    def reset_variable(self):
        if self.xp == np:  # below(31) of [1]
            self.Qx_power_FTM = self.xp.abs((self.Q_FMM[:, None] @ self.X_FTM[:, :, :, None])[:, :, :, 0]) ** 2
        else:
            self.Qx_power_FTM = self.xp.abs((self.Q_FMM[:, None] * self.X_FTM[:, :, None]).sum(axis=3)) ** 2
        self.Y_FTM = (self.lambda_NFT[..., None] * self.G_NM[:, None, None]).sum(axis=0)


    def initialize_covarianceMatrix(self):
        if "unit" in self.init_SCM:
            # tile相当于MATLAB的repmat Q_FMM是diagonalizer
            self.Q_FMM = self.xp.tile(self.xp.eye(self.n_mic), [self.n_freq, 1, 1]).astype(self.xp.complex)
            self.G_NM = self.xp.ones([self.n_source, self.n_mic], dtype=self.xp.float) * 1e-2
            for m in range(self.n_mic):
                self.G_NM[m % self.n_source, m] = 1  # diagonal initialization
        elif "obs" in self.init_SCM:
            mixture_covarianceMatrix_FMM = self.XX_FTMM.sum(axis=1) / (self.xp.trace(self.XX_FTMM, axis1=2, axis2=3).sum(axis=1))[:, None, None]
            eig_val, eig_vec = np.linalg.eigh(self.convert_to_NumpyArray(mixture_covarianceMatrix_FMM))  # 特征值特征向量
            self.Q_FMM = self.xp.asarray(eig_vec).transpose(0, 2, 1).conj()
            self.G_NM = self.xp.ones([self.n_source, self.n_mic], dtype=self.xp.float) / self.n_mic
            self.G_NM[0] = self.xp.asarray(eig_val[0])
        elif "ILRMA" in self.init_SCM:
            sys.path.append("../Rank1_Model")
            from ILRMA import ILRMA
            ilrma = ILRMA(n_basis=2, init_SCM="unit", xp=self.xp)
            ilrma.load_spectrogram(self.X_FTM)
            ilrma.solve(n_iteration=15, save_likelihood=False, save_wav=False, save_path="./", interval_save_parameter=1000)
            separated_spec_power = self.xp.abs(ilrma.separated_spec).mean(axis=(1, 2))
            self.Q_FMM = ilrma.SeparationMatrix_FMM
            self.G_NM = self.xp.ones([self.n_source, self.n_mic], dtype=self.xp.float) * 1e-2
            for n in range(self.n_source):
                self.G_NM[n, separated_spec_power.argmax()] = 1
                separated_spec_power[separated_spec_power.argmax()] = 0
        else:
            print("Please specify how to initialize covariance matrix {unit, obs}")
            raise ValueError

        self.normalize()


    def separate_FastWienerFilter(self, source_index=None, mic_index=MIC_INDEX):  # 要改
        Qx_FTM = (self.Q_FMM[:, None] * self.X_FTM[:, :, None]).sum(axis=3)
        if source_index != None:
            diagonalizer_inv_FMM = self.calculateInverseMatrix(self.Q_FMM)
            self.separated_spec = self.convert_to_NumpyArray((diagonalizer_inv_FMM[:, None] @ (Qx_FTM * ( (self.lambda_NFT[source_index, :, :, None] * self.G_NM[source_index, None, None]) / (self.lambda_NFT[..., None]* self.G_NM[:, None, None]).sum(axis=0) ) )[..., None])[:, :, mic_index, 0])
        else:
            for n in range(self.n_source):
                diagonalizer_inv_FMM = self.calculateInverseMatrix(self.Q_FMM)
                tmp = self.convert_to_NumpyArray((diagonalizer_inv_FMM[:, None] @ (Qx_FTM * ( (self.lambda_NFT[n, :, :, None] * self.G_NM[n, None, None]) / (self.lambda_NFT[..., None]* self.G_NM[:, None, None]).sum(axis=0) ) )[..., None])[:, :, mic_index, 0])
                if n == 0:
                    self.separated_spec = np.zeros([self.n_source, tmp.shape[0], tmp.shape[1]], dtype=np.complex)
                self.separated_spec[n] = tmp


    def save_parameter(self, fileName):
        param_list = [self.lambda_NFT, self.G_NM, self.Q_FMM]
        if self.xp != np:
            param_list = [self.convert_to_NumpyArray(param) for param in param_list]
        pic.dump(param_list, open(fileName, "wb"))


    def load_parameter(self, fileName):
        param_list = pic.load(open(fileName, "rb"))
        if self.xp != np:
            param_list = [cuda.to_gpu(param) for param in param_list]
        self.lambda_NFT, self.G_NM, self.Q_FMM = param_list

        self.n_source, self.n_freq, self.n_time = self.lambda_NFT.shape
        self.n_mic = self.G_NM.shape[-1]


class pre_batch_init_fast2():
    def __init__(self, a_W=None, b_W=None, W_NFK=None, G_a_1=None, G_b_1=None, G_NM=None):
        # W参数的传递
        self.a_W = a_W
        self.b_W = b_W
        self.W_NFK = W_NFK
        # G参数的传递
        self.G_a_1 = G_a_1
        self.G_b_1 = G_b_1
        self.G_NM = G_NM


if __name__ == "__main__":
    import argparse
    import pickle as pic
    import sys, os

    parser = argparse.ArgumentParser()
    parser.add_argument(    'input_fileName', type= str, help='filename of the multichannel observed signals')
    parser.add_argument(         '--file_id', type= str, default="None", help='file id')
    parser.add_argument(             '--gpu', type= int, default=     0, help='GPU ID')
    parser.add_argument(           '--n_fft', type= int, default=  1024, help='number of frequencies')
    parser.add_argument(        '--n_source', type= int, default=     2, help='number of noise')
    parser.add_argument(         '--n_basis', type= int, default=     8, help='number of basis')
    parser.add_argument(        '--init_SCM', type= str, default="ILRMA", help='unit, obs, ILRMA')
    parser.add_argument(     '--n_iteration', type= int, default=   100, help='number of iteration')
    parser.add_argument(       '--save_path', type= str, default="../output/", help='save path of wave file')
    parser.add_argument(        '--dl_ratio', type= float, default=     0, help='diagonal loading ratio')
    parser.add_argument(        '--save_wav', type= str, default=True, help='save wave file')
    parser.add_argument('--first_batch_size', type= int, default=40, help='frame num of first mini batch')
    parser.add_argument( '--mini_batch_size', type= int, default=4, help='frame num of each mini batch')
    parser.add_argument(      '--weight_rho', type= float, default=0.9, help='online update weight rho of last mini batch')
    args = parser.parse_args()

    if args.gpu < 0:
        import numpy as xp
    else:
        import cupy as xp
        print("Use GPU " + str(args.gpu))
        cuda.get_device_from_id(args.gpu).use()

    wav, fs = sf.read(args.input_fileName)
    wav = wav.T
    M = len(wav)
    for m in range(M):
        tmp = librosa.core.stft(wav[m], n_fft=args.n_fft, hop_length=int(args.n_fft/4))
        if m == 0:
            spec_all = np.zeros([tmp.shape[0], tmp.shape[1], M], dtype=np.complex)
        spec_all[:, :, m] = tmp

    T = spec_all.shape[1]
    first_batch_size = args.first_batch_size
    mini_batch_size = args.mini_batch_size
    all_sep_spec = 0

    for t in range(first_batch_size, T + 1):
        # online分帧
        if t == first_batch_size:
            tmp = spec_all[:, :t, :]
        else:
            tmp = spec_all[:, t - mini_batch_size:t, :]
        spec = np.zeros([tmp.shape[0], tmp.shape[1], tmp.shape[2]], dtype=np.complex)
        spec = tmp
        args.file_id = t + 1 - first_batch_size

        # 初始化separater对象
        separater = FastMNMF2_online(n_source=args.n_source, n_basis=args.n_basis, xp=xp, init_SCM=args.init_SCM, dl_ratio=args.dl_ratio, total_frame=T, all_sep_spec=all_sep_spec, first_batch_size=first_batch_size, weight_rho=args.weight_rho, mini_batch_size=args.mini_batch_size)

        # 初始化STFT频谱
        separater.load_spectrogram(spec)

        # 定义命名相关的参数
        separater.file_id = args.file_id
        separater.fs = fs

        # 定义pre_batch对象，用于存储上一batch的参数值
        if args.file_id == 1:
            pre_batch = pre_batch_init_fast2()
        else:
            # W参数的传递
            separater.pre_batch_a_W = pre_batch.a_W
            separater.pre_batch_b_W = pre_batch.b_W
            separater.pre_batch_W_NFK = pre_batch.W_NFK
            # G参数的传递
            separater.pre_batch_G_a_1 = pre_batch.G_a_1
            separater.pre_batch_G_b_1 = pre_batch.G_b_1
            separater.pre_batch_G_NM = pre_batch.G_NM

        # solve函数,负责算法核心运算、参数保存和结果输出
        all_sep_spec = separater.solve(n_iteration=args.n_iteration, save_likelihood=False, save_parameter=False, save_wav=args.save_wav, save_path=args.save_path, interval_save_parameter=25)

        # 在线的参数保存，提供给下一个mini-batch
        # W参数的保存
        pre_batch.a_W = separater.a_W
        pre_batch.b_W = separater.b_W
        pre_batch.W_NFK = separater.W_NFK
        # G参数的保存
        pre_batch.G_a_1 = separater.G_a_1
        pre_batch.G_b_1 = separater.G_b_1
        pre_batch.G_NM = separater.G_NM